#include <immintrin.h>
#include <math.h>
#include <stdio.h>
#include <time.h>
#include "utils.h"

void scalar_nested_conditions(const float *a, const float *b, const float *c,
                              const float *d, float *dst, int n)
{
    for (int i = 0; i < n; i++)
    {
        float sa = a[i];
        float sb = b[i];
        float sc = c[i];
        float sd = d[i];
        float final = 0;
        if (sa > 10)
        { // 第一层
            if (sb < 5)
            { // 第二层
                if (sc == 3)
                { // 第三层
                    if (sd != 0)
                    {                    // 第四层（新增）
                        final = sa + sb; // 加法
                    }
                    else
                    {
                        final = sa - sc; // 减法（操作数变化）
                    }
                }
                else
                {
                    final = sa * sb; // 乘法
                }
            }
            else
            {
                final = sa / sb; // 除法
            }
        }
        else
        {
            final = sqrt(sa); // 平方根（修改外层分支）
        }
        dst[i] = final;
    }
}

void simd_nested_conditions(const float *a, const float *b, const float *c,
                            const float *d, float *dst, int n)
{
    const __m256 zero = _mm256_setzero_ps();
    const __m256 three = _mm256_set1_ps(3.0);
    const __m256 ten = _mm256_set1_ps(10.0);
    const __m256 five = _mm256_set1_ps(5.0);
    for (int i = 0; i < n; i += 8)
    {
        // 加载数据
        __m256 va = _mm256_loadu_ps(a + i);
        __m256 vb = _mm256_loadu_ps(b + i);
        __m256 vc = _mm256_loadu_ps(c + i);
        __m256 vd = _mm256_loadu_ps(d + i);

        // 生成条件掩码（从外层到内层）
        __m256 mask1 = _mm256_cmp_ps(va, ten, _CMP_GT_OQ);   // a > 10
        __m256 mask2 = _mm256_cmp_ps(vb, five, _CMP_LT_OQ);  // b < 5
        __m256 mask3 = _mm256_cmp_ps(vc, three, _CMP_EQ_OQ); // c == 3
        __m256 mask4 = _mm256_cmp_ps(vd, zero, _CMP_NEQ_OQ); // d != 0

        // 预计算所有分支结果
        __m256 add = _mm256_add_ps(va, vb);
        __m256 sub = _mm256_sub_ps(va, vc); // 注意操作数变化
        __m256 mul = _mm256_mul_ps(va, vb);
        __m256 div = _mm256_div_ps(va, vb);
        __m256 sqrt_val = _mm256_sqrt_ps(va);

        __m256 cond1 = _mm256_and_ps(mask1, mask2);
        // 逐层混合（从最内层开始）
        // 第四层混合：mask1 & mask2 & mask3 & mask4
        __m256 cond4 = _mm256_and_ps(cond1, _mm256_and_ps(mask3, mask4));
        __m256 layer4 = _mm256_blendv_ps(sub, add, cond4);

        // 第三层混合：mask1 & mask2 & mask3
        __m256 cond3 = _mm256_and_ps(cond1, mask3);
        __m256 layer3 = _mm256_blendv_ps(mul, layer4, cond3);

        // 第二层混合：mask1 & mask2
        __m256 cond2 = cond1;
        __m256 layer2 = _mm256_blendv_ps(div, layer3, cond2);

        // 第一层混合：mask1
        __m256 final = _mm256_blendv_ps(sqrt_val, layer2, mask1);

        // 存储结果
        _mm256_storeu_ps(dst + i, final);
    }
}

int main()
{
    int len = 1;
    for (len = 1000; len < 100000000; len *= 10)
    {
        float *a = aligned_alloc(32, len * sizeof(float));
        float *b = aligned_alloc(32, len * sizeof(float));
        float *c = aligned_alloc(32, len * sizeof(float));
        float *d = aligned_alloc(32, len * sizeof(float));
        float *dst_scalar = aligned_alloc(32, len * sizeof(float));
        float *dst_simd = aligned_alloc(32, len * sizeof(float));

        // 初始化数据
        for (int i = 0; i < len; ++i)
        {
            a[i] = (float)rand() / RAND_MAX;
            b[i] = (float)rand() / RAND_MAX;
            c[i] = (float)rand() / RAND_MAX;
            d[i] = (float)rand() / RAND_MAX;
        }
        printf("数据个数 %d\n", len);
        // 标量版本计时
        BENCHMARK_FUNC(10, 100, scalar_nested_conditions, a, b, c, d, dst_scalar, len);

        // SIMD版本计时
        BENCHMARK_FUNC(10, 100, simd_nested_conditions, a, b, c, d, dst_simd, len);

        // 校验结果一致性
        for (int i = 0; i < len; ++i)
        {
            if (dst_scalar[i] != dst_simd[i])
            {
                printf("dst mismatch at index %d\n", i);
                break;
            }
        }

        free(a);
        free(b);
        free(c);
        free(d);

        free(dst_scalar);
        free(dst_simd);
    }
    return 0;
}